import platgo as pg


class DE(pg.Algorithm):

    type: dict = {'single': True, 'multi': False, 'many': False, 'real': True, 'binary': False, 'permutation': False,
                  "large": True, 'expensive': False, 'constrained': False, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}
    
    def __init__(self, maxgen: int, problem: pg.Problem) -> None:
        #  TODO 没有考虑约束
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = 'DE'
        self.mut = pg.operators.MutPol(problem)  # 多项式变异
        self.xov = pg.operators.DE(F=0.5)  # 差分算子
    
    def go(self, N: int = None, population: pg.Population = None):
        """
         main function for Different Evolution
         if population is None, generate a new population with N
        :param N: population size
        :param population: population to be optimized
        :return:
        """
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop()
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)  # 计算目标函数值

        while self.not_terminal(pop):
            offspring = self.xov(pop)
            offspring = self.mut(offspring)
            self.problem.cal_obj(offspring)  # 计算子代种群目标函数值
            temp_pop = offspring + pop  # 合并种群
            # 返回目标值小的一半
            idx = pg.utils.naive_selection(temp_pop, pop.N)
            pop = temp_pop[idx]
            # 目标函数值求和应该是逐步减小的
            # print("current generation: {}, the sum of objective: {}".format(self.gen, np.mean(pop.objv)))
        return pop
